import cv2
import numpy as np
import pandas as pd
import os
import glob

def extract_bar_chart_data(image_path: str, series_colors_rgb: list, y_axis_range: tuple, series_names: list, x_labels: list = None):
    """
    从分组条形图图像中自动提取数据。

    Args:
        image_path (str): 图表图像的文件路径。
        series_colors_rgb (list): 一个包含多个元组的列表，每个元组代表一个数据系列的RGB颜色。
                                  例如: [(255, 0, 0), (0, 255, 0)] 代表红色和绿色系列。
        y_axis_range (tuple): Y轴的数值范围，格式为 (最小值, 最大值)。例如: (0, 100)。
        series_names (list): 与颜色顺序对应的数据系列名称。
        x_labels (list): X轴标签列表，用于替换默认的"Group 1"、"Group 2"等。

    Returns:
        pandas.DataFrame: 包含提取数据的DataFrame，或在失败时返回None。
    """
    image = cv2.imread(image_path)
    if image is None:
        print(f"错误: 无法加载图片 '{image_path}'")
        return None

    # 创建一个用于调试的可视化副本
    output_image = image.copy()

    # --- 1. 自动检测坐标轴 ---
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)

    if lines is None:
        print("错误: 未能检测到足够长的直线作为坐标轴。")
        return None

    # 寻找最底部的水平线作为X轴，最左侧的垂直线作为Y轴
    x_axis_y_pixel, y_axis_x_pixel = -1, image.shape[1]
    y_axis_top_pixel = image.shape[0]

    for line in lines:
        x1, y1, x2, y2 = line[0]
        # 判断是水平线还是垂直线
        if abs(y1 - y2) < 5:  # 水平线
            if max(y1, y2) > x_axis_y_pixel:
                x_axis_y_pixel = max(y1, y2)
        elif abs(x1 - x2) < 5:  # 垂直线
            if min(x1, x2) < y_axis_x_pixel:
                y_axis_x_pixel = min(x1, x2)
                y_axis_top_pixel = min(y1, y2) # Y轴的顶端像素

    if x_axis_y_pixel == -1 or y_axis_x_pixel == image.shape[1]:
        print("错误: 未能清晰识别X轴或Y轴。")
        return None

    cv2.line(output_image, (0, x_axis_y_pixel), (image.shape[1], x_axis_y_pixel), (255, 0, 255), 2) # X轴(紫色)
    cv2.line(output_image, (y_axis_x_pixel, 0), (y_axis_x_pixel, image.shape[0]), (255, 0, 255), 2) # Y轴(紫色)

    # --- 2. 计算像素到数值的映射关系 ---
    y_min_val, y_max_val = y_axis_range
    yaxis_pixel_height = x_axis_y_pixel - y_axis_top_pixel
    value_range = y_max_val - y_min_val
    
    if yaxis_pixel_height <= 0:
        print("错误: Y轴像素高度计算错误。请检查坐标轴识别结果。")
        return None

    # --- 3. 按颜色提取每个系列的条形数据 ---
    all_series_data = {}
    color_tolerance = 25 # 颜色容差，以应对图片压缩等造成的颜色偏移

    for i, color_rgb in enumerate(series_colors_rgb):
        series_name = series_names[i]
        # OpenCV使用BGR格式，所以需要转换颜色
        color_bgr = np.array(color_rgb[::-1]) 
        
        # 创建颜色遮罩
        lower_bound = np.clip(color_bgr - color_tolerance, 0, 255)
        upper_bound = np.clip(color_bgr + color_tolerance, 0, 255)
        mask = cv2.inRange(image, lower_bound, upper_bound)

        # 寻找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        series_bars = []
        for contour in contours:
            # 过滤掉过小的噪点
            if cv2.contourArea(contour) < 100:
                continue

            x, y, w, h = cv2.boundingRect(contour)
            
            # --- 4. 计算数据 ---
            # 条形的高度（从Y轴顶端到底部Y坐标的距离）
            bar_pixel_height = x_axis_y_pixel - y
            
            # 将像素高度转换为真实数值
            value = (bar_pixel_height / yaxis_pixel_height) * value_range + y_min_val
            
            series_bars.append({'x_center': x + w / 2, 'value': value})
            
            # 在调试图像上绘制结果
            cv2.rectangle(output_image, (x, y), (x + w, y + h), (0, 255, 255), 2) # 黄色框
            cv2.putText(output_image, f"{value:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)

        # 按x轴位置排序，以确保分组顺序正确
        series_bars.sort(key=lambda item: item['x_center'])
        all_series_data[series_name] = [item['value'] for item in series_bars]

    # --- 5. 保存调试图像并整理输出 ---
    import os
    os.makedirs("data_chart/extract", exist_ok=True)
    
    # 生成基于输入文件名的调试图像文件名
    input_basename = os.path.splitext(os.path.basename(image_path))[0]
    debug_filename = f"data_chart/extract/{input_basename}_debug.png"
    cv2.imwrite(debug_filename, output_image)
    print(f"调试图像已保存为 '{debug_filename}'")

    # 将字典转换为Pandas DataFrame，处理数据长度不一致的情况
    try:
        # 找出所有系列中最长的数据长度
        max_length = max(len(values) for values in all_series_data.values()) if all_series_data else 0
        
        if max_length == 0:
            print("警告: 未检测到任何有效的条形数据")
            return None
        
        # 对每个系列进行长度标准化，缺失的数据用0补充
        standardized_data = {}
        for series_name, values in all_series_data.items():
            if len(values) < max_length:
                # 用0补充到最大长度
                padded_values = values + [0.0] * (max_length - len(values))
                print(f"系列 '{series_name}' 数据长度不足，已用0补充: {len(values)} -> {max_length}")
                standardized_data[series_name] = padded_values
            else:
                standardized_data[series_name] = values
        
        df = pd.DataFrame(standardized_data)
        # 使用用户提供的X轴标签，或使用默认的Group标签
        num_groups = max_length
        if x_labels and len(x_labels) >= num_groups:
            df.index = x_labels[:num_groups]
        else:
            df.index = [f"Group {i+1}" for i in range(num_groups)]
        
        # 保存提取的数据到data_chart/data文件夹，使用新的命名格式
        os.makedirs("data_chart/data", exist_ok=True)
        
        # 从文件名确定图表类型用于命名
        input_basename = os.path.splitext(os.path.basename(image_path))[0]
        if '_Task_' in input_basename or '_task_' in input_basename:
            chart_type_suffix = "Task"
        elif '_Node_' in input_basename or '_node_' in input_basename:
            chart_type_suffix = "Node"
        else:
            chart_type_suffix = "Task"  # 默认
        
        # 解析plotname (去掉ymin和ymax部分)
        parts = input_basename.split('_')
        if len(parts) >= 3:
            plotname = '_'.join(parts[:-2])  # 去掉最后两个部分(ymin, ymax)
        else:
            plotname = input_basename
            
        csv_filename = f"data_chart/data/{plotname}.csv"
        df.to_csv(csv_filename)
        print(f"提取的数据已保存为 '{csv_filename}'")
        
        return df
    except ValueError as e:
        print(f"错误: 无法创建DataFrame: {e}")
        print("检测到的数据:")
        for series_name, values in all_series_data.items():
            print(f"  {series_name}: {len(values)} 个值 - {values}")
        return None
    except Exception as e:
        print(f"创建DataFrame时发生未知错误: {e}")
        print("检测到的数据:", all_series_data)
        return None


def process_all_charts(input_dir: str, colors_rgb: list, series_names: list, task_x_labels: list, node_x_labels: list):
    """
    处理指定目录中的所有图表文件。
    
    Args:
        input_dir (str): 包含图表的目录路径
        colors_rgb (list): RGB颜色列表
        series_names (list): 数据系列名称列表
        task_x_labels (list): Task类型图表的X轴标签列表
        node_x_labels (list): Node类型图表的X轴标签列表
    
    Returns:
        dict: 包含所有图表提取结果的字典
    """
    # 支持的图片格式
    image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff']
    image_files = []
    
    # 收集所有图片文件
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(input_dir, ext)))
    
    if not image_files:
        print(f"错误: 在目录 '{input_dir}' 中未找到任何图片文件")
        return {}
    
    results = {}
    
    for image_path in image_files:
        filename = os.path.basename(image_path)
        name_without_ext = os.path.splitext(filename)[0]
        
        # 解析文件名格式: {plotname}_{ymin}_{ymax}
        parts = name_without_ext.split('_')
        if len(parts) < 3:
            print(f"警告: 文件 '{filename}' 格式不正确，跳过处理")
            print("期望格式: {plotname}_{ymin}_{ymax}")
            continue
        
        try:
            # 获取ymin和ymax（最后两个部分）
            ymax = float(parts[-1])
            ymin = float(parts[-2])
            # 剩余部分作为plotname
            plotname = '_'.join(parts[:-2])
            
            y_axis_range = (ymin, ymax)
            
            # 根据文件名确定使用哪种X轴标签
            if '_Task_' in name_without_ext or '_task_' in name_without_ext:
                x_labels = task_x_labels
                chart_type = "Task"
            elif '_Node_' in name_without_ext or '_node_' in name_without_ext:
                x_labels = node_x_labels
                chart_type = "Node"
            else:
                # 默认使用Task标签
                x_labels = task_x_labels
                chart_type = "Default(Task)"
            
            print(f"\n--- 处理图表: {filename} ---")
            print(f"图表名称: {plotname}")
            print(f"图表类型: {chart_type}")
            print(f"Y轴范围: {y_axis_range}")
            print(f"使用X轴标签: {x_labels}")
            
            # 提取数据
            extracted_df = extract_bar_chart_data(
                image_path, colors_rgb, y_axis_range, series_names, x_labels
            )
            
            if extracted_df is not None:
                results[plotname] = {
                    'dataframe': extracted_df,
                    'y_range': y_axis_range,
                    'filename': filename,
                    'chart_type': chart_type,
                    'x_labels_used': x_labels
                }
                print(f"成功提取数据:")
                print(extracted_df)
            else:
                print(f"提取数据失败")
                
        except ValueError as e:
            print(f"警告: 无法解析文件 '{filename}' 中的数值: {e}")
            continue
    
    return results


if __name__ == '__main__':

    # 输入目录
    input_directory = "data_chart/old_chart"
    
    # --- 配置通用参数 ---
    # 定义要识别的颜色 (16进制格式)
    COLORS_TO_FIND_HEX = [
        '#1f77b4',   #  - DE
        '#2CA02C',   #  - DOWN  
        '#9467BD',   #  - KUBE
        '#E377C2',   #  - PSO
        '#BCBD22',   #  - PPCM
        '#FF7F0E',   #  - DEP_SOFT
        '#D62728',   #  - GA
        '#8C564B',   #  - MONKEY
        '#7F7F7F',   #  - PPO
    ]
    
    # 将16进制转换为RGB格式供OpenCV使用
    def hex_to_rgb(hex_color):
        hex_color = hex_color.lstrip('#')
        return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    
    COLORS_TO_FIND_RGB = [hex_to_rgb(color) for color in COLORS_TO_FIND_HEX]
    SERIES_NAMES = ['DE', 'DOWN', 'KUBE','PSO','PPCM','DEP_SOFT','GA','MONKEY','PPO']
    
    # 定义两种类型的X轴标签
    TASK_X_LABELS = ['1500', '2000', '2500', '3000', '3500']  # Task number标签
    NODE_X_LABELS = ['9', '16', '25', '36']       # Node number标签
    
    print("--- 开始批量处理图表 ---")
    print(f"输入目录: {input_directory}")
    print(f"文件名格式要求: {{plotname}}_{{ymin}}_{{ymax}}")
    print(f"图表类型识别:")
    print(f"  - 包含'_Task_'或'_task_': 使用Task number标签 {TASK_X_LABELS}")
    print(f"  - 包含'_Node_'或'_node_': 使用Node number标签 {NODE_X_LABELS}")
    print(f"  - 其他情况: 默认使用Task number标签")
    
    # 处理所有图表
    all_results = process_all_charts(input_directory, COLORS_TO_FIND_RGB, SERIES_NAMES, TASK_X_LABELS, NODE_X_LABELS)
    
    if all_results:
        print(f"\n--- 批量处理完成 ---")
        print(f"成功处理了 {len(all_results)} 个图表:")
        for plotname, result in all_results.items():
            print(f"- {plotname}: Y轴范围{result['y_range']}, 类型: {result['chart_type']}")
            
        # 可选：保存所有结果到一个汇总文件
        os.makedirs("data_chart", exist_ok=True)
        summary_data = {}
        for plotname, result in all_results.items():
            summary_data[plotname] = {
                'data': result['dataframe'].to_dict(),
                'y_range': result['y_range'],
                'chart_type': result['chart_type'],
                'x_labels_used': result['x_labels_used']
            }
        
        import json
        with open("data_chart/all_extracted_data_summary.json", 'w', encoding='utf-8') as f:
            json.dump(summary_data, f, indent=2, ensure_ascii=False)
        print("\n所有提取结果已保存到 'data_chart/all_extracted_data_summary.json'")
    else:
        print("\n--- 未能成功处理任何图表 ---")
